#include "njm_swapchain.hpp"

#include <array>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <limits>
#include <set>
#include <stdexcept>
#include <algorithm>


namespace njm {
        
    NjmSwapChain::NjmSwapChain(NjmDevice &deviceRef, VkExtent2D windowExtent) 
        : device{deviceRef}, windowExtent{windowExtent}
    { init(); }

    NjmSwapChain::NjmSwapChain(NjmDevice &deviceRef, VkExtent2D windowExtent, std::shared_ptr<NjmSwapChain> previous) 
        : device{deviceRef}, windowExtent{windowExtent}, oldSwapChain{previous}
    {
        init();
        oldSwapChain = nullptr;
    }

    void NjmSwapChain::init()
    {
        createSwapChain();
        createImageViews();
        createRenderPass();
        createDepthResources();
        createFramebuffers();
        createSyncObjects();
    }
    
    NjmSwapChain::~NjmSwapChain()
    {
        for (auto &&imgView : swapChainImageViews)
        {
            if(imgView) vkDestroyImageView(device.device(), imgView, nullptr);
        }
        swapChainImageViews.clear();

        for (auto &&depthView : depthImageViews)
        {
            if(depthView) vkDestroyImageView(device.device(), depthView, nullptr);
        }
        depthImageViews.clear();

        for (auto &&depthImg : depthImages)
        {
            if(depthImg) vkDestroyImage(device.device(), depthImg, nullptr);            
        }
        depthImages.clear();
        
        for (auto &&mem : depthImageMemorys)
        {
            if(mem) vkFreeMemory(device.device(), mem, nullptr);
        }
        depthImageMemorys.clear();

        for (auto &&imgSema : imageAvailableSemaphores)
        {
            if(imgSema) vkDestroySemaphore(device.device(), imgSema, nullptr);
        }
        imageAvailableSemaphores.clear();

        for (auto &&renderSema : renderFinishedSemaphores)
        {
            if(renderSema) vkDestroySemaphore(device.device(), renderSema, nullptr);
        }
        renderFinishedSemaphores.clear();

        for (auto &&flightFence : inFlightFences)
        {
            if(flightFence) vkDestroyFence(device.device(), flightFence, nullptr);
        }
        inFlightFences.clear();
        // for (auto &&imgFence : imagesInFlight)
        // {
        //     if(imgFence) vkDestroyFence(device.device(), imgFence, nullptr);
        // }
        // imagesInFlight.clear();
        
        for (auto &&framebuf : swapChainFramebuffers)
        {
            if(framebuf) vkDestroyFramebuffer(device.device(), framebuf, nullptr);
        }
        swapChainFramebuffers.clear();

        if(renderPass) 
        {
            vkDestroyRenderPass(device.device(), renderPass, nullptr);
            renderPass = nullptr;
        }
        if (swapChain != nullptr)
        { 
            vkDestroySwapchainKHR(device.device(), swapChain, nullptr);
            swapChain = nullptr;  
        }
        // TODO
    }

    // creator
    void NjmSwapChain::createSwapChain()
    {     
        VkSwapchainCreateInfoKHR scInfo{};
        scInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
        scInfo.surface = device.getWindowSurface();  
        scInfo.minImageCount = 2;
        scInfo.imageFormat = VK_FORMAT_B8G8R8A8_SRGB;
        scInfo.imageColorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR;
        scInfo.imageExtent = windowExtent;
        scInfo.imageArrayLayers = 1;
        scInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
        scInfo.preTransform = VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR;
        scInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
        scInfo.presentMode = VK_PRESENT_MODE_FIFO_KHR;
        scInfo.clipped = VK_TRUE;
        scInfo.oldSwapchain = oldSwapChain == nullptr ? VK_NULL_HANDLE : oldSwapChain->swapChain;

        if (vkCreateSwapchainKHR(device.device(), &scInfo, nullptr, &swapChain) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to create swapchain");
        }
        swapChainImageFormat = scInfo.imageFormat;
        swapChainExtent = scInfo.imageExtent;
    }

    void NjmSwapChain::createImageViews()
    {
        uint32_t imgCnt;
        vkGetSwapchainImagesKHR(device.device(), swapChain, &imgCnt, nullptr);
        if (imgCnt == 0) {
            throw std::runtime_error("swapchain image count is zero!");
        }
        swapChainImages.resize(imgCnt);
        vkGetSwapchainImagesKHR(device.device(), swapChain, &imgCnt, swapChainImages.data());

        swapChainImageViews.resize(imgCnt);
        for (size_t i = 0; i < imgCnt; i++)
        {
            VkComponentMapping compMap{};
            compMap.r = VK_COMPONENT_SWIZZLE_IDENTITY;
            compMap.g = VK_COMPONENT_SWIZZLE_IDENTITY;
            compMap.b = VK_COMPONENT_SWIZZLE_IDENTITY;
            compMap.a = VK_COMPONENT_SWIZZLE_IDENTITY;
            VkImageSubresourceRange subResRange{};
            subResRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            subResRange.baseArrayLayer = 0;
            subResRange.layerCount = 1;
            subResRange.baseMipLevel = 0;
            subResRange.levelCount = 1;

            VkImageViewCreateInfo ivInfo{};
            ivInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
            ivInfo.format = swapChainImageFormat;
            ivInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
            ivInfo.image = swapChainImages[i];
            ivInfo.components = compMap;
            ivInfo.subresourceRange = subResRange;

            if (vkCreateImageView(device.device(), &ivInfo, nullptr, &swapChainImageViews[i]) != VK_SUCCESS)
            { throw std::runtime_error("failed to create imageview[i]"); }
        }       
    }

    void NjmSwapChain::createRenderPass()
    {
        VkAttachmentDescription colorAttachDescrip{};
        colorAttachDescrip.format = swapChainImageFormat;
        colorAttachDescrip.samples = VK_SAMPLE_COUNT_1_BIT;
        colorAttachDescrip.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
        colorAttachDescrip.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
        colorAttachDescrip.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
        colorAttachDescrip.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
        colorAttachDescrip.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        colorAttachDescrip.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;

        VkAttachmentDescription depthStencilAttachDescrip{};
        // depthStencilAttachDescrip.format = swapChainDepthFormat;
        depthStencilAttachDescrip.format = VK_FORMAT_D32_SFLOAT_S8_UINT;
        depthStencilAttachDescrip.samples = VK_SAMPLE_COUNT_1_BIT;
        depthStencilAttachDescrip.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
        depthStencilAttachDescrip.storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
        depthStencilAttachDescrip.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
        depthStencilAttachDescrip.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
        depthStencilAttachDescrip.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        depthStencilAttachDescrip.finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

        std::array<VkAttachmentDescription, 2> attachs;
        attachs[0] = colorAttachDescrip;
        attachs[1] = depthStencilAttachDescrip;

        VkAttachmentReference colorAttachRef{};
        colorAttachRef.attachment = 0;
        colorAttachRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
        VkAttachmentReference depthStencilAttachRef{};
        depthStencilAttachRef.attachment = 1;
        depthStencilAttachRef.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

        VkSubpassDescription subpass{};
        subpass.colorAttachmentCount = 1;
        subpass.pColorAttachments = &colorAttachRef;
        subpass.pDepthStencilAttachment = &depthStencilAttachRef;
        subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
        // subpass.inputAttachmentCount = ;
        // subpass.pInputAttachments = ;

        VkSubpassDependency deps{};
        deps.srcSubpass = VK_SUBPASS_EXTERNAL;
        deps.srcAccessMask = 0;
        deps.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
        deps.dstSubpass = 0;
        deps.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
        deps.dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
        // VkSubpassDependency dependency = {};
        // dependency.srcSubpass = VK_SUBPASS_EXTERNAL;
        // dependency.dstSubpass = 0;
        // dependency.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
        // dependency.dstStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
        // dependency.srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
        // dependency.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;

        VkRenderPassCreateInfo rpInfo{};
        rpInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
        rpInfo.attachmentCount = static_cast<uint32_t>(attachs.size());
        rpInfo.pAttachments = attachs.data();
        rpInfo.subpassCount = 1;
        rpInfo.pSubpasses = &subpass;
        rpInfo.dependencyCount = 1;
        rpInfo.pDependencies = &deps;
        if (vkCreateRenderPass(device.device(), &rpInfo, nullptr, &renderPass) != VK_SUCCESS)
        { throw std::runtime_error("failed to create renderpass"); }
    }

    void NjmSwapChain::createDepthResources()
    {
        // VkFormat depthFormat = findDepthFormat();
        VkFormat depthFormat = VK_FORMAT_D32_SFLOAT_S8_UINT;
        swapChainDepthFormat = depthFormat;
        VkExtent2D swapChainExtent = getSwapChainExtent();

        depthImages.resize(imageCount());
        depthImageMemorys.resize(imageCount());
        depthImageViews.resize(imageCount());

        for (size_t i = 0; i < depthImages.size(); i++)
        {
            VkImageCreateInfo depImgInfo{};
            depImgInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
            depImgInfo.imageType = VK_IMAGE_TYPE_2D;
            depImgInfo.extent.width = swapChainExtent.width;
            depImgInfo.extent.height = swapChainExtent.height;
            depImgInfo.extent.depth = 1;
            depImgInfo.mipLevels = 1;
            depImgInfo.arrayLayers = 1;
            depImgInfo.format = depthFormat;
            depImgInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
            depImgInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
            depImgInfo.usage = VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT;
            depImgInfo.samples = VK_SAMPLE_COUNT_1_BIT;
            depImgInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
            depImgInfo.flags = 0;
            device.createImageWithInfo(depImgInfo, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, depthImages[i], depthImageMemorys[i]);

            VkComponentMapping compMap{};
            compMap.r = VK_COMPONENT_SWIZZLE_IDENTITY;
            compMap.g = VK_COMPONENT_SWIZZLE_IDENTITY;
            compMap.b = VK_COMPONENT_SWIZZLE_IDENTITY;
            compMap.a = VK_COMPONENT_SWIZZLE_IDENTITY;
            VkImageViewCreateInfo imgViewInfo{};
            imgViewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
            imgViewInfo.image = depthImages[i];
            imgViewInfo.format = depthFormat;
            imgViewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
            imgViewInfo.components = compMap;
            imgViewInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;
            imgViewInfo.subresourceRange.baseMipLevel = 0;
            imgViewInfo.subresourceRange.levelCount = 1;
            imgViewInfo.subresourceRange.baseArrayLayer = 0;
            imgViewInfo.subresourceRange.layerCount = 1;
            if (vkCreateImageView(device.device(), &imgViewInfo, nullptr, &depthImageViews[i]) != VK_SUCCESS) {
                throw std::runtime_error("can't create imgview");
            }
        }       
    }

    void NjmSwapChain::createFramebuffers()
    {
        uint32_t fbCnt = imageCount();
        swapChainFramebuffers.resize(fbCnt);
        for (size_t i = 0; i < fbCnt; i++)
        {
            std::array<VkImageView, 2> attachs = {swapChainImageViews[i], depthImageViews[i]};
            VkExtent2D ex = getSwapChainExtent();

            VkFramebufferCreateInfo fbInfo{};
            fbInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
            fbInfo.renderPass = renderPass;
            fbInfo.attachmentCount = static_cast<uint32_t>(attachs.size());
            fbInfo.pAttachments = attachs.data();
            fbInfo.width = ex.width;
            fbInfo.height = ex.height;
            fbInfo.layers = 1;
            if (vkCreateFramebuffer(device.device(), &fbInfo, nullptr, &swapChainFramebuffers[i]) != VK_SUCCESS)
            { throw std::runtime_error("failed to create framebuffer[i]"); }
        }
    }

    void NjmSwapChain::createSyncObjects()
    {
        imageAvailableSemaphores.resize(MAX_FRAMES_IN_FLIGHT);
        renderFinishedSemaphores.resize(MAX_FRAMES_IN_FLIGHT);
        inFlightFences.resize(MAX_FRAMES_IN_FLIGHT);
        imagesInFlight.resize(imageCount(), VK_NULL_HANDLE);

        VkSemaphoreCreateInfo semaInfo{};
        semaInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;
        VkFenceCreateInfo fenceInfo{};
        fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
        fenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;

        for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
        {
            if (vkCreateSemaphore(device.device(), &semaInfo, nullptr, &imageAvailableSemaphores[i]) != VK_SUCCESS)
            { throw std::runtime_error("failed to create imageAvailableSemaphores[i]"); }
            if (vkCreateSemaphore(device.device(), &semaInfo, nullptr, &renderFinishedSemaphores[i]) != VK_SUCCESS)
            { throw std::runtime_error("failed to create renderFinishedSemaphores[i]"); }
            if (vkCreateFence(device.device(), &fenceInfo, nullptr, &inFlightFences[i]) != VK_SUCCESS)
            { throw std::runtime_error("failed to create inFlightFences[i]"); }   
        }
    }

    // helper
    VkPresentModeKHR NjmSwapChain::chooseSwapPresentMode(const std::vector<VkPresentModeKHR> &availableModes) 
    {
        for (const auto &mode : availableModes) {
            if (mode == VK_PRESENT_MODE_MAILBOX_KHR) {
                std::cout << "vk-swapchain-present-mode : mailbox" << std::endl;
                return mode;
            }
        }
        std::cout << "vk-swapchain-present-mode : fifo(v-sync)" << std::endl;
        return VK_PRESENT_MODE_FIFO_KHR;
    }
    VkSurfaceFormatKHR NjmSwapChain::chooseSwapSurfaceFormat(const std::vector<VkSurfaceFormatKHR> &availableFormats)
    {
        for (const auto &format : availableFormats) {
            if (format.format == VK_FORMAT_B8G8R8A8_SRGB &&
                format.colorSpace == VK_COLORSPACE_SRGB_NONLINEAR_KHR) {
                std::cout << "vk-swapchain-surface-format.format : VK_FORMAT_B8G8R8A8_SRGB" << std::endl;
                std::cout << "vk-swapchain-surface-format.colorSpace : VK_COLORSPACE_SRGB_NONLINEAR_KHR" << std::endl;
                return format;
            }            
        }
        return availableFormats[0];
    }
    VkExtent2D NjmSwapChain::chooseSwapExtent(const VkSurfaceCapabilitiesKHR &capabilities)
    {
        if (capabilities.currentExtent.width != std::numeric_limits<uint32_t>::max()) {
            return capabilities.currentExtent; // 如果 currentExtent 的宽度和高度不是最大值，直接使用 currentExtent
        } else {
            // int width, height; // 获取窗口的帧缓冲大小（以像素为单位）
            // glfwGetFramebufferSize(window, &width, &height);
            // // 将窗口大小转换为 VkExtent2D 类型
            // VkExtent2D actualExtent = {
            //     static_cast<uint32_t>(width),
            //     static_cast<uint32_t>(height)
            // };
            VkExtent2D actualExtent{};
            actualExtent.width = std::clamp(windowExtent.width, capabilities.minImageExtent.width, capabilities.maxImageExtent.width);
            actualExtent.height = std::clamp(windowExtent.height, capabilities.minImageExtent.height, capabilities.maxImageExtent.height);
            return actualExtent;
        }
    }
    VkFormat NjmSwapChain::findSupportedFormat(const std::vector<VkFormat>& candidates, VkImageTiling tiling, VkFormatFeatureFlags features) 
    {
        for (VkFormat format : candidates) {
            VkFormatProperties props;
            vkGetPhysicalDeviceFormatProperties(device.getPhysicalDevice(), format, &props);
            if (tiling == VK_IMAGE_TILING_LINEAR && (props.linearTilingFeatures & features) == features) {
                return format;
            } else if (tiling == VK_IMAGE_TILING_OPTIMAL && (props.optimalTilingFeatures & features) == features) {
                return format;
            }
        }
        throw std::runtime_error("failed to find supported format!");
    }

    // public api
    VkFormat NjmSwapChain::findDepthFormat()
    {
        return findSupportedFormat(
            {VK_FORMAT_D32_SFLOAT, VK_FORMAT_D32_SFLOAT_S8_UINT, VK_FORMAT_D24_UNORM_S8_UINT},
            VK_IMAGE_TILING_OPTIMAL, VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT); 
    }

    VkResult NjmSwapChain::aquireNextImage(uint32_t *imageIndex)
    {
        auto rst = vkWaitForFences(device.device(), 1, &inFlightFences[currentFrame], VK_TRUE, std::numeric_limits<uint64_t>::max());
        if (rst!= VK_SUCCESS){throw std::runtime_error("failed to wait for fence!");}
        vkResetFences(device.device(), 1, &inFlightFences[currentFrame]);
        return vkAcquireNextImageKHR(device.device(), swapChain, std::numeric_limits<uint64_t>::max(), 
                    imageAvailableSemaphores[currentFrame], VK_NULL_HANDLE, imageIndex);
    }

    VkResult NjmSwapChain::submitCommandBuffers(const VkCommandBuffer *buffers, uint32_t *imageIndex)
    {
        // TODO
        imagesInFlight[*imageIndex] = inFlightFences[currentFrame];        
        VkQueue q;
        vkGetDeviceQueue(device.device(), device.getGraphicsQueueFamilyIndex(), 0, &q);
        
        VkSubmitInfo subInfo{};
        subInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
        subInfo.commandBufferCount = 1;
        subInfo.pCommandBuffers = buffers;
        subInfo.waitSemaphoreCount = 1;
        subInfo.pWaitSemaphores = &imageAvailableSemaphores[currentFrame];
        subInfo.signalSemaphoreCount = 1;
        subInfo.pSignalSemaphores = &renderFinishedSemaphores[currentFrame];
        VkPipelineStageFlags waitStageFlags[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
        subInfo.pWaitDstStageMask = waitStageFlags;

        auto rst = vkQueueSubmit(q, 1, &subInfo, imagesInFlight[*imageIndex]);
        if (rst != VK_SUCCESS){throw std::runtime_error("failed to submit cmdbuf!");}
        // present
        VkPresentInfoKHR presentInfo{};
        presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
        presentInfo.swapchainCount = 1;
        presentInfo.pSwapchains = &swapChain;
        presentInfo.waitSemaphoreCount = 1;
        presentInfo.pWaitSemaphores = &renderFinishedSemaphores[currentFrame];
        presentInfo.pImageIndices = imageIndex;
        presentInfo.pResults = nullptr;
        rst = vkQueuePresentKHR(q, &presentInfo);
        return rst;
    }
    
}